package com.shujia.flink.sink

import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.{RichSinkFunction, SinkFunction}
import org.apache.flink.streaming.api.scala._

import java.sql.{Connection, DriverManager, PreparedStatement}

object Demo4MysqlSInk {
  def main(args: Array[String]): Unit = {
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    env.setParallelism(2)
    /**
     * 统计单词的数量，将统计结果保存到mysql中
     *
     */

    val linesDS: DataStream[String] = env.socketTextStream("master", 8888)

    //统计单词的数量
    val countDS: DataStream[(String, Int)] = linesDS
      .flatMap(_.split(","))
      .map((_, 1))
      .keyBy(_._1)
      .sum(1)

    //使用自定义sink
    countDS.addSink(new MysqlSink())

    env.execute()

  }
}

/**
 * 自定义sink
 * 1、实现SinkFunction接口
 * 2、重写invoke方法
 * 3、在invoke方法中将数据保存到外部数据库中
 *
 * SinkFunction
 * RichSinkFunction,比SinkFunction多了open和close方法，可以用于创建和关闭连接
 * open和close是在一个task中执行一次（相当于spark的foreachPartition，每一个创建创建一个连接）
 *
 */
class MysqlSink extends RichSinkFunction[(String, Int)] {

  //成员变量，可以在所有的方法中获取
  var con: Connection = _

  /**
   * open:每一个task中执行一次，task启动时执行
   */
  override def open(parameters: Configuration): Unit = {
    println("创建数据库连接")
    //1、加载驱动
    Class.forName("com.mysql.jdbc.Driver")
    //2、创建数据量连接
    con = DriverManager.getConnection("jdbc:mysql://master:3306/student", "root", "123456")
  }

  /**
   * 每一个task中执行一次,task取消时执行
   */
  override def close(): Unit = {
    //6、关闭连接
    con.close()
  }

  /**
   * invoke:每一条数据执行一次
   *
   * @param kv      ：每一条数据
   * @param context ：上下文对象
   */
  override def invoke(kv: (String, Int), context: SinkFunction.Context): Unit = {

    //3、写sql，插入数据
    //replace into: 替换插入，如果主键存在就替换，如果不存在就插入，穿点表时需要增加主键
    val stat: PreparedStatement = con.prepareStatement("replace into word_count(word,num) values(?,?)")
    //4、设置值
    stat.setString(1, kv._1)
    stat.setInt(2, kv._2)
    //5、执行插入
    stat.execute()

  }
}